import torch
import torch.nn as nn
import torch.optim as optim
from LeNet数据集 import *
from static.LeNet import LeNet
from torch.optim.lr_scheduler import ReduceLROnPlateau

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = LeNet().to(device)
Epochs = 100
lr = 0.001
loss_fn = nn.CrossEntropyLoss()

# 更换为RMSprop优化器，通常对某些问题有更好的表现
optimizer = optim.RMSprop(model.parameters(), lr=lr, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)

# 学习率调度器，使用ReduceLROnPlateau在验证损失停滞时降低学习率
scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5, verbose=True)

# 评估模型函数
def evaluate_model(model, val_loader, loss_fn, device):
    model.eval()
    val_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = loss_fn(outputs, labels)
            val_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    val_loss /= len(val_loader.dataset)
    accuracy = 100. * correct / total
    return val_loss, accuracy

best_loss = float('inf')

# 循环训练
for epoch in range(Epochs):
    loss_sum = 0
    model.train()
    for i, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        # 清空梯度
        optimizer.zero_grad()
        # 前向传播
        y_pred = model(images)
        # 计算损失
        loss = loss_fn(y_pred, labels)
        loss_sum += loss.item()
        # 反向传播
        loss.backward()
        # 更新参数
        optimizer.step()
    # 计算平均损失
    avg_loss = loss_sum / len(train_loader)
    # 在验证集上评估模型
    val_loss, val_accuracy = evaluate_model(model, val_loader, loss_fn, device)
    # 学习率调度器更新
    scheduler.step(val_loss)
    print(f"Epoch {epoch + 1} 训练损失: {avg_loss:.4f}, 验证损失: {val_loss:.4f}, 验证准确率: {val_accuracy:.2f}%")

    if val_loss < best_loss:
        best_loss = val_loss
        # 保存最佳模型
        torch.save(model.state_dict(), '../static/LeNet_mammals_best_model.pth')
        print(f"已保存最佳模型 (验证损失: {val_loss:.4f})")

# 保存最终训练的模型
torch.save(model.state_dict(), 'static/LeNet_mammals_model.pth')